import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import torch

def rodrigues_rotation(v1, v2):
    # Normalize the input vectors
    v1 = v1 / torch.norm(v1)
    v2 = v2 / torch.norm(v2)
    
    # Compute the cross product of the vectors
    k = torch.cross(v1, v2)

    # Compute the cosine of the angle between the vectors
    cos_theta = torch.dot(v1, v2)

    # Compute the sine of the angle between the vectors
    sin_theta = torch.norm(k)

    # Compute the skew-symmetric matrix using the cross product vector
    kx = torch.tensor([[0, -k[2], k[1]],
                      [k[2], 0, -k[0]],
                      [-k[1], k[0], 0]])

    # Compute the rotation matrix using Rodrigues' rotation formula
    R = torch.eye(3) + kx + torch.mm(kx, kx) * ((1 - cos_theta) / sin_theta**2) if sin_theta != 0 else torch.eye(3)

    return R

def cov_to_ellipsoid(V):
        # Compute eigenvalues and eigenvectors of V
    eigenvalues, eigenvectors = np.linalg.eig(V)

    # Generate points on a unit sphere
    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, np.pi, 100)
    x = np.outer(np.cos(u), np.sin(v))
    y = np.outer(np.sin(u), np.sin(v))
    z = np.outer(np.ones_like(u), np.cos(v))
    points = np.array([x.flatten(), y.flatten(), z.flatten()])

    # Transform the points using eigenvectors and eigenvalues
    transformed_points = (eigenvectors @ np.diag(np.sqrt(eigenvalues)) @ points).T
    
    return transformed_points

scale=10.

pc_B=torch.rand([3,128])*scale
pc=pc_B
pc_range=torch.sqrt(torch.sum(pc**2,axis=0))


_V=torch.diag(torch.tensor([scale*scale,scale,scale]))

unitX=torch.tensor([1.,0.,0.])


# Plot the ellipsoid
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')



ax.quiver(0,0,0,scale,0,0,color='red')
ax.quiver(0,0,0,0,scale,0,color='green')
ax.quiver(0,0,0,0,0,scale,color='blue')


for i in range(10):

    R=rodrigues_rotation(unitX,pc_B[:,i])
    d=torch.norm(pc_B[:,i])

    V=R@_V@R.T*d*d/1000

    transformed_points=cov_to_ellipsoid(V)

    transformed_points_1=pc_B[:,i]+transformed_points


    ax.quiver(0,0,0,pc_B[0,i],pc_B[1,i],pc_B[2,i],color='orange')
    transformed_points=transformed_points_1
    ax.plot_surface(transformed_points[:, 0].reshape([100,100]),
                    transformed_points[:, 1].reshape([100,100]),
                    transformed_points[:, 2].reshape([100,100]),
                    alpha=0.5)



ax.set_xlim(-0, scale)
ax.set_ylim(-0, scale)
ax.set_zlim(-0, scale)
# Set axis labels and title
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.title('Ellipsoid Visualization')

# Show the plot
plt.show()